"""
    本程序通过调用Google OR-Tools工具包，通过加压的方式求解满表的车站作业计划，进而求解高速铁路车站的通过能力
    作者：廖正文 北京交通大学

    运行本程序前，需要安装Google OR-Tools的Python包，安装方法如下：
    保持计算机在可以访问互联网的状态，在命令提示符中，输入
    python -m pip install --user ortools
    即可完成安装

    如果安装过程较慢，可以尝试使用清华大学TUNA的镜像，即输入
    python -m pip install --user ortools -i https://pypi.tuna.tsinghua.edu.cn/simple
"""

from ortools.linear_solver import pywraplp  # 载入 Google OR-Tools 求解工具包
import PyStationCapacityData as pscdata
import os


### 全局变量 ###

Data = None
''' 数据 '''

M = 100000
''' 大M '''

FixedVarDict = dict()
''' 已知变量字典 '''

# AppliedSolverName = "CP-SAT"  # 设置CP-SAT作为线性规划的求解器
AppliedSolverName = "Gurobi"  # 设置Gurobi作为线性规划的求解器
# AppliedSolverName = "SCIP"  # 设置SCIP作为线性规划的求解器
''' 采用求解器名称 '''

ObjFunc = "Mind"
# ObjFunc = "Mind-a"

ObjVal = -1
''' 目标函数值 '''

MIPGap = -1
''' 求解Gap '''

Solver = None
''' 求解器 '''

VarDict = dict()
''' 决策变量字典 '''

### 全局变量结束 ###


def Initialize():
    ''' 初始化 '''
    global Solver
    
    FixedVarDict.clear()
    VarDict.clear()
    Solver = pywraplp.Solver.CreateSolver(AppliedSolverName)

    Data.GenerateCandidateRoute()  # 生成可行路径
    Data.GenerateCandidateCell()  # 生成轨道电路

    Solver.SetTimeLimit(3600000)  # 最大求解时间限制，毫秒
    Solver.SetNumThreads(12)  # 使用线程数量

    if (os.path.exists(Data.OutputFileDirectory) != 1):  # 创建输出文件目录
        os.mkdir(Data.OutputFileDirectory)


def BuildModel():
    ''' 创建模型 '''
    Solver.Clear()
    
    BuildVariables()  # 构建决策变量
    BuildConstraints()  # 构建约束条件
    BuildObjective()  # 构建目标函数


def BuildVariables():
    ''' 创建决策变量 '''
    
    VarDict.clear()

    for train in Data.TrainList:
        for m in train.TrainMovementList:
            # 作业m的开始时间
            a_VarName = "a_{0}".format(m.ID)
            a = Solver.NumVar(0, 86400, a_VarName)
            VarDict[a_VarName] = a

            # 作业m的结束时间
            d_VarName = "d_{0}".format(m.ID)
            d = Solver.NumVar(0, 86400, d_VarName)
            VarDict[d_VarName] = d

            for r in m.CandidateRouteList:
                # 作业m的进路选择
                x_VarName = "x_{0}_{1}".format(m.ID, r.ID)
                x = Solver.BoolVar(x_VarName)
                VarDict[x_VarName] = x

            for c in m.CandidateCellList:
                # 作业占用轨道电路开始时间
                s_VarName = "s_{0}_{1}".format(m.ID, c.ID)
                s = Solver.NumVar(0, 86400, s_VarName)
                VarDict[s_VarName] = s

                # 作业占用轨道电路结束时间
                e_VarName = "e_{0}_{1}".format(m.ID, c.ID)
                e = Solver.NumVar(0, 86400, e_VarName)
                VarDict[e_VarName] = e

    for tr1 in Data.TrainList:
        for tr2 in Data.TrainList:
            if (tr1 == tr2):
                continue
            for m1 in tr1.TrainMovementList:
                for m2 in tr2.TrainMovementList:
                    for c in Data.CurrentStation.CellList:
                        if c in m1.CandidateCellList and c in m2.CandidateCellList:
                            # 作业占用先后次序
                            y_VarName = "y_{0}_{1}_{2}".format(c.ID, m1.ID, m2.ID)
                            if y_VarName not in FixedVarDict:  # 如果变量y的值已知，则不创建
                                y = Solver.BoolVar(y_VarName)
                                VarDict[y_VarName] = y


def BuildObjective():
    ''' 创建目标函数 '''
    objExpr = 0

    if ObjFunc == "Mind":
        for train in Data.TrainList:
            for m in train.TrainMovementList:
                # 作业m的结束时间
                d_VarName = "d_{0}".format(m.ID)
                d = VarDict[d_VarName]
                objExpr += d

    elif ObjFunc == "Mind-a":
        for train in Data.TrainList:
            for m in train.TrainMovementList:
                # 作业m的结束时间
                d_VarName = "d_{0}".format(m.ID)
                d = VarDict[d_VarName]
                a_VarName = "a_{0}".format(m.ID)
                a = VarDict[a_VarName]
                objExpr += d-a
    elif ObjFunc == "Mind+a":
        for train in Data.TrainList:
            for m in train.TrainMovementList:
                # 作业m的结束时间
                d_VarName = "d_{0}".format(m.ID)
                d = VarDict[d_VarName]
                a_VarName = "a_{0}".format(m.ID)
                a = VarDict[a_VarName]
                objExpr += (d+a)
                
    Solver.Minimize(objExpr)  # 目标函数为最小化所有作业的结束时间之和


def BuildConstraints():
    ''' 创建约束条件 '''
    BuildMovementDurationConstraint()  # 创建作业持续时间约束
    BuildTrainConservationConstraint()  # 创建列车作业时间连续性约束
    BuildTrainRouteConnConstraint()  # 创建列车作业空间连续性约束
    BuildMovementTimeWindowConstraint()  # 创建作业时间窗约束
    BuildRouteSelectionConstraint()  # 创建作业进路选择约束
    BuildCellOccupationTimeConstraint()  # 创建轨道电路占用时间约束
    BuildCellOccupationConflictConstraint()  # 创建轨道电路占用冲突约束
    BuildFixedVarConstrain()  # 创建固定变量约束


def BuildFixedVarConstrain():
    ''' 创建固定变量约束 '''
    for train in Data.TrainList:
        for m in train.TrainMovementList:
            for r in m.CandidateRouteList:
                x_VarName = "x_{0}_{1}".format(m.ID, r.ID)
                if x_VarName in FixedVarDict:
                    x = VarDict[x_VarName]
                    Solver.Add(x == FixedVarDict[x_VarName],"RtFixedCt_{0}_{1}".format(m.ID, r.ID))   # 创建固定径路选择变量


def BuildMovementDurationConstraint():
    ''' 创建作业持续时间约束 '''
    for tr in Data.TrainList:
        for m in tr.TrainMovementList:
            a = VarDict["a_{0}".format(m.ID)]
            d = VarDict["d_{0}".format(m.ID)]

            Solver.Add(d - a >= m.MinDuration, "MinDurCt_{0}".format(m.ID))
            Solver.Add(d - a <= m.MaxDuration, "MaxDurCt_{0}".format(m.ID))


def BuildTrainConservationConstraint():
    ''' 创建列车作业时间连续性约束 '''
    for tr in Data.TrainList:
        lastMovement = None
        for m in tr.TrainMovementList:
            if lastMovement == None:
                lastMovement = m
                continue
            last_d = VarDict["d_{0}".format(lastMovement.ID)]
            current_a = VarDict["a_{0}".format(m.ID)]
            Solver.Add(last_d == current_a,
                       "TrainConsCt_{0}_{1}".format(lastMovement.ID, m.ID))
            lastMovement = m


def BuildTrainRouteConnConstraint():
    ''' 创建列车作业空间连续性约束 '''
    for tr in Data.TrainList:
        m1 = None
        for m2 in tr.TrainMovementList:
            if m1 == None:
                m1 = m2
                continue

            for r1 in m1.CandidateRouteList:
                x1 = VarDict["x_{0}_{1}".format(m1.ID, r1.ID)]
                sum_x2_expr = 0
                for r2 in m2.CandidateRouteList:
                    x2 = VarDict["x_{0}_{1}".format(m2.ID, r2.ID)]
                    if not Data.IsRouteWithConn(r1,r2):
                        sum_x2_expr += x2
                Solver.Add(sum_x2_expr + (1-x1)*M>= 0, "TrainNoConn1Ct_{0}_{1}".format(m1.ID, r1.ID))
                Solver.Add(sum_x2_expr <= (1-x1)*M, "TrainNoConn2Ct_{0}_{1}".format(m1.ID, r1.ID))
            m1 = m2


def BuildMovementTimeWindowConstraint():
    ''' 创建作业时间窗约束 '''
    for tr in Data.TrainList:
        for m in tr.TrainMovementList:
            a = VarDict["a_{0}".format(m.ID)]
            Solver.Add(a >= Data.ConvertDateTimeToTotalSec(
                m.OrgStartTime) - Data.LeftShiftTolerance, "aWinLCt_{0}".format(m.ID))
            Solver.Add(a <= Data.ConvertDateTimeToTotalSec(
                m.OrgStartTime) + Data.RightShiftTolerance, "aWinRCt_{0}".format(m.ID))

            d = VarDict["d_{0}".format(m.ID)]
            Solver.Add(d >= Data.ConvertDateTimeToTotalSec(
                m.OrgEndTime) - Data.LeftShiftTolerance, "dWinLCt_{0}".format(m.ID))
            Solver.Add(d <= Data.ConvertDateTimeToTotalSec(
                m.OrgEndTime) + Data.RightShiftTolerance, "dWinRCt_{0}".format(m.ID))


def BuildRouteSelectionConstraint():
    ''' 创建作业进路选择约束 '''
    for tr in Data.TrainList:
        for m in tr.TrainMovementList:
            sum_x = 0
            for r in m.CandidateRouteList:
                x_VarName = "x_{0}_{1}".format(m.ID, r.ID)                    
                x = VarDict[x_VarName]
                sum_x += x
            Solver.Add(sum_x == 1, "RtSelCt_{0}".format(m.ID))


def BuildCellOccupationTimeConstraint():
    ''' 创建轨道电路占用时间约束 '''
    for tr in Data.TrainList:
        for m in tr.TrainMovementList:
            # 作业m的开始时间
            a_VarName = "a_{0}".format(m.ID)
            a = VarDict[a_VarName]

            # 作业m的结束时间
            d_VarName = "d_{0}".format(m.ID)
            d = VarDict[d_VarName]

            for r in m.CandidateRouteList:
                x_VarName = "x_{0}_{1}".format(m.ID, r.ID)
                x = VarDict[x_VarName]

                for c in r.CellList:
                    # 作业占用轨道电路开始时间
                    s_VarName = "s_{0}_{1}".format(m.ID, c.ID)
                    s = VarDict[s_VarName]

                    # 作业占用轨道电路结束时间
                    e_VarName = "e_{0}_{1}".format(m.ID, c.ID)
                    e = VarDict[e_VarName]

                    cot = Data.CurrentStation.CellOccupationTimeSetList[0].CellOccupationTimeDict[r,c]
                    startShift = cot.StartOccupationShift
                    endShift = cot.EndOccupationShift

                    Solver.Add(s + M * (1-x) >= a + startShift,
                               "cotCts1_{0}_{1}".format(m.ID, c.ID))
                    Solver.Add(s <= a + startShift + M * (1-x),
                               "cotCts2_{0}_{1}".format(m.ID, c.ID))
                    Solver.Add(e + M * (1-x) >= d + endShift,
                               "cotCte1_{0}_{1}".format(m.ID, c.ID))
                    Solver.Add(e <= d + endShift + M * (1-x),
                               "cotCte2_{0}_{1}".format(m.ID, c.ID))


def BuildCellOccupationConflictConstraint():
    ''' 创建轨道电路占用冲突约束 '''
    for tr1 in Data.TrainList:
        for tr2 in Data.TrainList:
            if tr1 == tr2:
                continue

            for m1 in tr1.TrainMovementList:
                for m2 in tr2.TrainMovementList:
                    for c in Data.CurrentStation.CellList:
                        if c in m1.CandidateCellList and c in m2.CandidateCellList:
                            # m2作业占用轨道电路开始时间
                            s_VarName = "s_{0}_{1}".format(m2.ID, c.ID)
                            s_m2 = VarDict[s_VarName]

                            # m1作业占用轨道电路结束时间
                            e_VarName = "e_{0}_{1}".format(m1.ID, c.ID)
                            e_m1 = VarDict[e_VarName]

                            y_VarName = "y_{0}_{1}_{2}".format(c.ID, m1.ID, m2.ID)

                            if y_VarName in FixedVarDict:  # 顺序固定
                                if FixedVarDict[y_VarName] == 1:
                                    Solver.Add(s_m2 >= e_m1,
                                        "ConfCt_{0}_{1}_{2}".format(c.ID, m1.ID, m2.ID))
                            else:  # 顺序不固定
                                y = VarDict[y_VarName]
                                Solver.Add((1-y)*M + s_m2 >= e_m1,
                                    "ConfCt_{0}_{1}_{2}".format(c.ID, m1.ID, m2.ID))
                            
                                # y的关系
                                yVerse_VarName = "y_{0}_{1}_{2}".format(c.ID, m2.ID, m1.ID)
                                yVerse = VarDict[yVerse_VarName]
                                Solver.Add(y+yVerse == 1, "yCt_{0}_{1}_{2}".format(c.ID, m1.ID, m2.ID))


def WriteLpFile(fileName):
    ''' 
        输出*.lp文件字符串，并保存至指定文件中
        lp格式的文件是线性（整数）规划模型完全展开后的形式，可以根据此文件检查以上生成模型的代码是否存在问题
    '''
    lp_str = Solver.ExportModelAsLpFormat(False)
    lp_file = open(fileName, "w")
    lp_file.writelines(lp_str)
    lp_file.close()


def ParseSolution():
    ''' 自求解器中解析求解结果，生成压缩后的列车运行图数据 '''
    global ObjVal
    global Solver
    global MIPGap

    # 目标函数值
    ObjVal = Solver.Objective().Value()
    bestBound = Solver.Objective().BestBound()
    MIPGap = (ObjVal-bestBound)/ObjVal

    # 列车运行计划
    for tr in Data.TrainList:
        for m in tr.TrainMovementList:
            a = VarDict["a_{0}".format(m.ID)]
            d = VarDict["d_{0}".format(m.ID)]
            # print("{0},{1},{2},{3}".format(tr.ID, m.ID, a.SolutionValue(), d.SolutionValue()))
            m.StartTime = Data.ConvertTotalSecToDateTime(a.SolutionValue())
            m.EndTime = Data.ConvertTotalSecToDateTime(d.SolutionValue())

            for r in m.CandidateRouteList:
                x_VarName = "x_{0}_{1}".format(m.ID, r.ID)
                if x_VarName not in FixedVarDict:
                    x = VarDict[x_VarName]
                    if x.SolutionValue() > 0.99:
                        FixedVarDict[x_VarName] = 1
                        m.SelectingRoute = r
                    else:
                        FixedVarDict[x_VarName] = 0
                elif FixedVarDict[x_VarName] == 1:
                    m.SelectingRoute = r

            for c in m.CandidateCellList:
                # 作业占用轨道电路开始时间
                s_VarName = "s_{0}_{1}".format(m.ID, c.ID)
                s = VarDict[s_VarName]
                sVal = s.SolutionValue()

                # 作业占用轨道电路结束时间
                e_VarName = "e_{0}_{1}".format(m.ID, c.ID)
                e = VarDict[e_VarName]
                eVal = e.SolutionValue()

                m.OccupancyTimeDict[c] = (sVal, eVal)
    
    # 列车运行顺序
    for tr1 in Data.TrainList:
        for tr2 in Data.TrainList:
            if (tr1 == tr2):
                continue
            for m1 in tr1.TrainMovementList:
                for m2 in tr2.TrainMovementList:
                    for c in Data.CurrentStation.CellList:
                        if c in m1.CandidateCellList and c in m2.CandidateCellList:
                            # 作业占用先后次序
                            y_VarName = "y_{0}_{1}_{2}".format(c.ID, m1.ID, m2.ID)
                            if y_VarName not in FixedVarDict:
                                y = VarDict[y_VarName]

                                if y.SolutionValue() > 0.99:
                                    FixedVarDict[y_VarName] = 1
                                else:
                                    FixedVarDict[y_VarName] = 0
    print("求解数据解析完成！")


def Solve():
    """ 求解车站能力计算模型 """
    global ObjVal
    ObjVal = -1

    Solver.EnableOutput()  # 允许求解引擎向终端输出日志
    WriteLpFile(Data.OutputFileDirectory +  "/PyStationCapacityModel.lp")  # 输出模型lp文件

    status = Solver.Solve()

    if  status == pywraplp.Solver.OPTIMAL or status == pywraplp.Solver.FEASIBLE:
        ParseSolution()
        print("求解结束！")
        return True
    return False
